#include "burner.h"
#include "png.h"

#define SSHOT_NOERROR 0
#define SSHOT_ERROR_BPP_NOTSUPPORTED 1
#define SSHOT_LIBPNG_ERROR 2
#define SSHOT_OTHER_ERROR 3
#define MAX_PREV_NUM 20

static unsigned char* convertedImage = NULL;

// ==> save preview image, modified by regret
static void PreviewName(wchar_t* szName)
{
	wchar_t path[MAX_PATH] = L"";

	wcscpy(path, getMiscPath(PATH_PREVIEW));
	// get md image path
	if ((BurnDrvGetHardwareCode() & HARDWARE_PUBLIC_MASK) == HARDWARE_SEGA_MEGADRIVE) {
		wcscat(path, L"MD\\");
	}

	swprintf(szName, L"%s%hs.png", path, BurnDrvGetTextA(DRV_NAME));
	if (!fileExists(szName)) {
		return;
	}

	for (int i = 2; i < MAX_PREV_NUM; i++) {
		swprintf(szName, L"%s%hs [%02i].png", path, BurnDrvGetTextA(DRV_NAME), i);
		if (!fileExists(szName))
			return;
	}
}
// <== save preview image

// convert image, added and fixed rotation by regret
static inline unsigned char* ConvertVidImage(const unsigned char* src, int flipVertical)
{
	const unsigned char* image = src;
	unsigned char* temp = NULL;
    unsigned w, h;

	BurnDrvGetVisibleSize(w, h);

	// Rotate and flip the image
	if (BurnDrvGetFlags() & BDF_ORIENTATION_VERTICAL) {
		unsigned char* oldTemp = temp;
		temp = new unsigned char[w * h * sizeof(int)];

		for (int x = 0; x < h; x++) {
			for (int y = 0; y < w; y++) {
				if (BurnDrvGetFlags() & BDF_ORIENTATION_FLIPPED) {
					((unsigned int*)temp)[(w - y - 1) + x * w] = ((unsigned int*)image)[x + y * h];
				} else {
					((unsigned int*)temp)[y + (h - x - 1) * w] = ((unsigned int*)image)[x + y * h];
				}
			}
		}

		delete [] oldTemp;
		convertedImage = temp;
		image = convertedImage;
	}
	else if (BurnDrvGetFlags() & BDF_ORIENTATION_FLIPPED) {
		unsigned char* oldTemp = temp;
		temp = new unsigned char[w * h * sizeof(int)];

		for (int y = h - 1; y >= 0; y--) {
			for (int x = w - 1; x >= 0; x--) {
				((unsigned int*)temp)[(w - x - 1) + (h - y - 1) * w] = ((unsigned int*)image)[x + y * w];
			}
		}

		delete [] oldTemp;
		convertedImage = temp;
		image = convertedImage;
	}

	if (flipVertical) {
		unsigned char* oldTemp = temp;
		temp = new unsigned char[w * h * sizeof(int)];

		for (int y = 0; y < h; y++) {
			for (int x = 0; x < w; x++) {
				((unsigned int*)temp)[x + (h - y - 1) * w] = ((unsigned int*)image)[x + y * w];
			}
		}

		delete [] oldTemp;
		convertedImage = temp;
		image = convertedImage;
	}

	return convertedImage;
}

int makeScreenShot(bool screenShot)
{
	char szAuthor[MAX_PATH]; char szDescription[MAX_PATH]; char szCopyright[MAX_PATH]; char szSoftware[MAX_PATH]; char szSource[MAX_PATH];
	png_text text_ptr[8] = { { 0, 0, 0, 0 }, };
	int num_text = 8;

	time_t currentTime;
	tm* tmTime;
	png_time_struct png_time;

	wchar_t sShotName[MAX_PATH];
	unsigned w, h;

	// do our PNG construct things
	png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	if (!png_ptr) {
		return SSHOT_LIBPNG_ERROR;
	}

	png_infop info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr) {
		png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		return SSHOT_LIBPNG_ERROR;
	}

	if (setjmp(png_jmpbuf(png_ptr))) {
		png_destroy_write_struct(&png_ptr, &info_ptr);
		_wremove(sShotName);
		return SSHOT_LIBPNG_ERROR;
	}

	if (!videoInterface.getVidImage()) {
		return SSHOT_OTHER_ERROR;
	}

	BurnDrvGetVisibleSize(w, h);

	// Convert the image to 32-bit
	unsigned char* pSShot = ConvertVidImage(videoInterface.getVidImage(), 0);
	if (!pSShot) {
		pSShot = videoInterface.getVidImage();
	}

	// Get the time
	time(&currentTime);
	tmTime = localtime(&currentTime);
	png_convert_from_time_t(&png_time, currentTime);

	// create dir if dir doesn't exist
	if (!directoryExists(getMiscPath(PATH_SCREENSHOT))) {
		CreateDirectory(getMiscPath(PATH_SCREENSHOT), NULL);
	}

	if (screenShot) {
		// construct our filename -> "romname-mm-dd-hms.png"
		swprintf(sShotName, L"%s%hs-%.2d-%.2d-%.2d%.2d%.2d.png",
			getMiscPath(PATH_SCREENSHOT), BurnDrvGetTextA(DRV_NAME),
			tmTime->tm_mon + 1, tmTime->tm_mday, tmTime->tm_hour, tmTime->tm_min, tmTime->tm_sec);
	} else {
		// Make preview
		PreviewName(sShotName);
	}

	FILE* file = _wfopen(sShotName, L"wb");
	if (file == NULL) {
		png_destroy_write_struct(&png_ptr, &info_ptr);

		delete [] convertedImage;
		convertedImage = NULL;

		return SSHOT_OTHER_ERROR;
	}

	// Fill the PNG text fields
	sprintf(szAuthor, "%ls v%.20ls", APP_TITLE, appBurnVer);
	sprintf(szDescription, "Screenshot of %s", decorateGameName(BurnGetDrvSelect()));
	sprintf(szCopyright, "%s %s", BurnDrvGetTextA(DRV_DATE), BurnDrvGetTextA(DRV_MANUFACTURER));
	sprintf(szSoftware, "%ls v%.20ls using LibPNG ", APP_TITLE, PNG_LIBPNG_VER_STRING, appBurnVer);
	sprintf(szSource, "%s video game hardware", BurnDrvGetTextA(DRV_SYSTEM));

	text_ptr[0].key = "Title";			text_ptr[0].text = (char*)BurnDrvGetTextA(DRV_FULLNAME);
	text_ptr[1].key = "Author";			text_ptr[1].text = szAuthor;
	text_ptr[2].key = "Description";	text_ptr[2].text = szDescription;
	text_ptr[3].key = "Copyright";		text_ptr[3].text = szCopyright;
	text_ptr[4].key = "Creation Time";	text_ptr[4].text = png_convert_to_rfc1123(png_ptr, &png_time);
	text_ptr[5].key = "Software";		text_ptr[5].text = szSoftware;
	text_ptr[6].key = "Source";			text_ptr[6].text = szSource;
	text_ptr[7].key = "Comment";		text_ptr[7].text = "This screenshot was created by running the game in an emulator; it might not accurately reflect the actual hardware the game was designed to run on.";

	for (int i = 0; i < num_text; i++) {
		text_ptr[i].compression = PNG_TEXT_COMPRESSION_NONE;
	}

	png_set_text(png_ptr, info_ptr, text_ptr, num_text);

	png_init_io(png_ptr, file);

	png_set_IHDR(png_ptr, info_ptr, w, h, 8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
	png_write_info(png_ptr, info_ptr);

	png_set_filler(png_ptr, 0, PNG_FILLER_AFTER);

	png_set_bgr(png_ptr);

	png_bytep* pSShotImageRows = new png_bytep[h];
	for (int y = 0; y < h; y++) {
		pSShotImageRows[y] = pSShot + (y * w * sizeof(int));
	}

	png_write_image(png_ptr, pSShotImageRows);
	png_write_end(png_ptr, info_ptr);

	delete [] pSShotImageRows;
	fclose(file);
	png_destroy_write_struct(&png_ptr, &info_ptr);

	delete [] convertedImage;
	convertedImage = NULL;

	return SSHOT_NOERROR;
}
